// Copyright (c) 2024 IndustrialTBB Framework
// Licensed under MIT License

#ifndef ITBB_CORE_TBB_OBSERVER_H_
#define ITBB_CORE_TBB_OBSERVER_H_

#include <oneapi/tbb.h>
#include <oneapi/tbb/info.h>
#include <oneapi/tbb/task_scheduler_observer.h>

#include <atomic>
#include <chrono>
#include <functional>
#include <memory>
#include <unordered_map>
#include <vector>

namespace itbb {
// TBB线程监控数据
struct ThreadMetrics {
  std::atomic<uint64_t> tasks_executed{0};
  std::atomic<uint64_t> total_execution_time_ns{0};
  std::atomic<uint64_t> idle_time_ns{0};
  std::atomic<uint64_t> context_switches{0};
  std::chrono::steady_clock::time_point thread_start_time;
  std::chrono::steady_clock::time_point last_activity_time;

  ThreadMetrics()
      : thread_start_time(std::chrono::steady_clock::now()),
        last_activity_time(std::chrono::steady_clock::now()) {}
};

// TBB任务执行监控
class TbbTaskExecutionObserver : public tbb::task_scheduler_observer {
 public:
  explicit TbbTaskExecutionObserver(bool local_observation = false)
      : tbb::task_scheduler_observer(local_observation),
        total_worker_threads_(0),
        active_tasks_(0),
        peak_active_tasks_(0) {
    observe(true);
  }

  ~TbbTaskExecutionObserver() { observe(false); }

  // TBB回调实现
  void on_scheduler_entry(bool is_worker) override {
    if (is_worker) {
      total_worker_threads_.fetch_add(1, std::memory_order_relaxed);

      // 为每个工作线程初始化指标
      thread_local static bool metrics_initialized = false;
      if (!metrics_initialized) {
        InitializeThreadMetrics();
        metrics_initialized = true;
      }
    }

    // 记录线程进入时间
    thread_local auto entry_time = std::chrono::high_resolution_clock::now();
    GetThreadMetrics().last_activity_time = std::chrono::steady_clock::now();
  }

  void on_scheduler_exit(bool is_worker) override {
    if (is_worker) {
      total_worker_threads_.fetch_sub(1, std::memory_order_relaxed);
    }

    // 计算线程执行时间
    thread_local auto exit_time = std::chrono::high_resolution_clock::now();
    thread_local auto entry_time = std::chrono::high_resolution_clock::now();

    auto execution_duration =
        std::chrono::duration_cast<std::chrono::nanoseconds>(exit_time -
                                                             entry_time);

    auto& metrics = GetThreadMetrics();
    metrics.total_execution_time_ns.fetch_add(execution_duration.count(),
                                              std::memory_order_relaxed);
  }

  // 任务开始执行回调
  void on_task_begin() {
    active_tasks_.fetch_add(1, std::memory_order_relaxed);

    // 更新峰值活跃任务数
    uint64_t current_active = active_tasks_.load(std::memory_order_acquire);
    uint64_t current_peak = peak_active_tasks_.load(std::memory_order_acquire);

    while (current_active > current_peak &&
           !peak_active_tasks_.compare_exchange_weak(
               current_peak, current_active, std::memory_order_release)) {
      current_peak = peak_active_tasks_.load(std::memory_order_acquire);
    }

    // 记录任务开始时间
    thread_local auto task_start_time =
        std::chrono::high_resolution_clock::now();

    GetThreadMetrics().last_activity_time = std::chrono::steady_clock::now();
  }

  // 任务结束执行回调
  void on_task_end() {
    active_tasks_.fetch_sub(1, std::memory_order_relaxed);

    // 计算任务执行时间
    thread_local auto task_end_time = std::chrono::high_resolution_clock::now();
    thread_local auto task_start_time =
        std::chrono::high_resolution_clock::now();

    auto task_duration = std::chrono::duration_cast<std::chrono::nanoseconds>(
        task_end_time - task_start_time);

    auto& metrics = GetThreadMetrics();
    metrics.tasks_executed.fetch_add(1, std::memory_order_relaxed);
    metrics.total_execution_time_ns.fetch_add(task_duration.count(),
                                              std::memory_order_relaxed);
    metrics.last_activity_time = std::chrono::steady_clock::now();
  }

  // 获取总体统计信息
  struct SystemMetrics {
    uint64_t total_worker_threads;
    uint64_t active_tasks;
    uint64_t peak_active_tasks;
    uint64_t total_tasks_executed;
    uint64_t total_execution_time_ns;
    double average_task_duration_ns;
    double thread_utilization;
    size_t hardware_concurrency;
    size_t numa_nodes;
  };

  SystemMetrics GetSystemMetrics() const {
    SystemMetrics metrics{};

    metrics.total_worker_threads =
        total_worker_threads_.load(std::memory_order_acquire);
    metrics.active_tasks = active_tasks_.load(std::memory_order_acquire);
    metrics.peak_active_tasks =
        peak_active_tasks_.load(std::memory_order_acquire);
    metrics.hardware_concurrency = tbb::info::default_concurrency();

    // 计算所有线程的总统计信息
    uint64_t total_tasks = 0;
    uint64_t total_time = 0;

    for (const auto& [thread_id, thread_metrics] : thread_metrics_map_) {
      total_tasks +=
          thread_metrics->tasks_executed.load(std::memory_order_acquire);
      total_time += thread_metrics->total_execution_time_ns.load(
          std::memory_order_acquire);
    }

    metrics.total_tasks_executed = total_tasks;
    metrics.total_execution_time_ns = total_time;
    metrics.average_task_duration_ns =
        total_tasks > 0 ? static_cast<double>(total_time) / total_tasks : 0.0;

    // 计算线程利用率
    metrics.thread_utilization = CalculateThreadUtilization();

    // NUMA信息（简化实现）
    metrics.numa_nodes = 1;  // 可以通过系统调用获取实际NUMA节点数

    return metrics;
  }

  // 获取每个线程的详细指标
  std::vector<ThreadMetrics> GetPerThreadMetrics() const {
    std::vector<ThreadMetrics> result;
    result.reserve(thread_metrics_map_.size());

    for (const auto& [thread_id, metrics] : thread_metrics_map_) {
      result.push_back(*metrics);
    }

    return result;
  }

  // 重置统计信息
  void ResetMetrics() {
    total_worker_threads_.store(0, std::memory_order_release);
    active_tasks_.store(0, std::memory_order_release);
    peak_active_tasks_.store(0, std::memory_order_release);

    for (auto& [thread_id, metrics] : thread_metrics_map_) {
      metrics->tasks_executed.store(0, std::memory_order_release);
      metrics->total_execution_time_ns.store(0, std::memory_order_release);
      metrics->idle_time_ns.store(0, std::memory_order_release);
      metrics->context_switches.store(0, std::memory_order_release);
      metrics->thread_start_time = std::chrono::steady_clock::now();
      metrics->last_activity_time = std::chrono::steady_clock::now();
    }
  }

  // 设置事件回调
  using TaskEventCallback = std::function<void(const std::string&, uint64_t)>;
  void SetTaskEventCallback(TaskEventCallback callback) {
    task_event_callback_ = std::move(callback);
  }

  using ThreadEventCallback = std::function<void(std::thread::id, bool)>;
  void SetThreadEventCallback(ThreadEventCallback callback) {
    thread_event_callback_ = std::move(callback);
  }

 private:
  // 原子统计变量
  std::atomic<uint64_t> total_worker_threads_;
  std::atomic<uint64_t> active_tasks_;
  std::atomic<uint64_t> peak_active_tasks_;

  // 每线程指标存储
  tbb::concurrent_hash_map<std::thread::id, std::unique_ptr<ThreadMetrics>>
      thread_metrics_map_;

  // 事件回调
  TaskEventCallback task_event_callback_;
  ThreadEventCallback thread_event_callback_;

  // 线程本地指标初始化
  void InitializeThreadMetrics() {
    auto thread_id = std::this_thread::get_id();

    tbb::concurrent_hash_map<std::thread::id,
                             std::unique_ptr<ThreadMetrics>>::accessor accessor;
    if (thread_metrics_map_.insert(accessor, thread_id)) {
      accessor->second = std::make_unique<ThreadMetrics>();
    }

    // 触发线程事件回调
    if (thread_event_callback_) {
      thread_event_callback_(thread_id, true);
    }
  }

  // 获取当前线程的指标对象
  ThreadMetrics& GetThreadMetrics() {
    auto thread_id = std::this_thread::get_id();

    tbb::concurrent_hash_map<std::thread::id,
                             std::unique_ptr<ThreadMetrics>>::accessor accessor;
    if (thread_metrics_map_.find(accessor, thread_id)) {
      return *accessor->second;
    }

    // 如果找不到，创建新的（理论上不应该发生）
    if (thread_metrics_map_.insert(accessor, thread_id)) {
      accessor->second = std::make_unique<ThreadMetrics>();
    }

    return *accessor->second;
  }

  // 计算线程利用率
  double CalculateThreadUtilization() const {
    auto now = std::chrono::steady_clock::now();
    uint64_t total_possible_time = 0;
    uint64_t total_active_time = 0;

    for (const auto& [thread_id, metrics] : thread_metrics_map_) {
      auto thread_lifetime =
          std::chrono::duration_cast<std::chrono::nanoseconds>(
              now - metrics->thread_start_time)
              .count();

      total_possible_time += thread_lifetime;
      total_active_time +=
          metrics->total_execution_time_ns.load(std::memory_order_acquire);
    }

    return total_possible_time > 0
               ? static_cast<double>(total_active_time) / total_possible_time
               : 0.0;
  }
};

// TBB内存使用监控观察器
class TbbMemoryObserver : public tbb::task_scheduler_observer {
 public:
  struct MemoryMetrics {
    std::atomic<uint64_t> total_allocated_bytes{0};
    std::atomic<uint64_t> peak_allocated_bytes{0};
    std::atomic<uint64_t> current_allocated_bytes{0};
    std::atomic<uint64_t> allocation_count{0};
    std::atomic<uint64_t> deallocation_count{0};
    std::atomic<uint64_t> failed_allocations{0};
  };

  TbbMemoryObserver() : tbb::task_scheduler_observer(false) { observe(true); }

  ~TbbMemoryObserver() { observe(false); }

  void on_scheduler_entry(bool is_worker) override {
    // 记录线程进入时的内存使用情况
    if (is_worker) {
      RecordMemorySnapshot();
    }
  }

  void on_scheduler_exit(bool is_worker) override {
    // 记录线程退出时的内存使用情况
    if (is_worker) {
      RecordMemorySnapshot();
    }
  }

  // 记录内存分配
  void RecordAllocation(size_t bytes) {
    memory_metrics_.total_allocated_bytes.fetch_add(bytes,
                                                    std::memory_order_relaxed);
    memory_metrics_.current_allocated_bytes.fetch_add(
        bytes, std::memory_order_relaxed);
    memory_metrics_.allocation_count.fetch_add(1, std::memory_order_relaxed);

    // 更新峰值内存使用
    uint64_t current =
        memory_metrics_.current_allocated_bytes.load(std::memory_order_acquire);
    uint64_t peak =
        memory_metrics_.peak_allocated_bytes.load(std::memory_order_acquire);

    while (current > peak &&
           !memory_metrics_.peak_allocated_bytes.compare_exchange_weak(
               peak, current, std::memory_order_release)) {
      peak =
          memory_metrics_.peak_allocated_bytes.load(std::memory_order_acquire);
    }
  }

  // 记录内存释放
  void RecordDeallocation(size_t bytes) {
    memory_metrics_.current_allocated_bytes.fetch_sub(
        bytes, std::memory_order_relaxed);
    memory_metrics_.deallocation_count.fetch_add(1, std::memory_order_relaxed);
  }

  // 记录分配失败
  void RecordAllocationFailure() {
    memory_metrics_.failed_allocations.fetch_add(1, std::memory_order_relaxed);
  }

  // 获取内存统计信息
  MemoryMetrics GetMemoryMetrics() const { return memory_metrics_; }

  // 获取系统内存信息
  struct SystemMemoryInfo {
    uint64_t total_physical_memory;
    uint64_t available_physical_memory;
    uint64_t process_memory_usage;
    uint64_t tbb_memory_usage;
    double memory_fragmentation_ratio;
  };

  SystemMemoryInfo GetSystemMemoryInfo() const {
    SystemMemoryInfo info{};

    // 获取系统内存信息（平台相关实现）
    info.total_physical_memory = GetTotalPhysicalMemory();
    info.available_physical_memory = GetAvailablePhysicalMemory();
    info.process_memory_usage = GetProcessMemoryUsage();
    info.tbb_memory_usage =
        memory_metrics_.current_allocated_bytes.load(std::memory_order_acquire);

    // 计算内存碎片率
    info.memory_fragmentation_ratio = CalculateMemoryFragmentation();

    return info;
  }

 private:
  MemoryMetrics memory_metrics_;

  void RecordMemorySnapshot() {
    // 记录当前内存快照
    // 可以集成到性能监控系统
  }

  uint64_t GetTotalPhysicalMemory() const {
    // 平台相关的系统内存获取
    return 0;  // 简化实现
  }

  uint64_t GetAvailablePhysicalMemory() const {
    // 平台相关的可用内存获取
    return 0;  // 简化实现
  }

  uint64_t GetProcessMemoryUsage() const {
    // 平台相关的进程内存使用获取
    return 0;  // 简化实现
  }

  double CalculateMemoryFragmentation() const {
    // 计算内存碎片率
    return 0.0;  // 简化实现
  }
};

// TBB负载均衡监控观察器
class TbbLoadBalanceObserver : public tbb::task_scheduler_observer {
 public:
  struct LoadBalanceMetrics {
    std::atomic<uint64_t> task_steals{0};
    std::atomic<uint64_t> failed_steals{0};
    std::atomic<uint64_t> work_distribution_imbalance{0};
    std::vector<std::atomic<uint64_t>> per_thread_task_counts;

    LoadBalanceMetrics() {
      size_t thread_count = tbb::info::default_concurrency();
      per_thread_task_counts.resize(thread_count);
    }

    double GetStealSuccessRate() const {
      uint64_t total_steals = task_steals.load(std::memory_order_acquire);
      uint64_t failed = failed_steals.load(std::memory_order_acquire);

      return total_steals > 0
                 ? static_cast<double>(total_steals - failed) / total_steals
                 : 0.0;
    }

    double GetLoadImbalanceRatio() const {
      if (per_thread_task_counts.empty()) return 0.0;

      uint64_t min_tasks = UINT64_MAX;
      uint64_t max_tasks = 0;

      for (const auto& count : per_thread_task_counts) {
        uint64_t tasks = count.load(std::memory_order_acquire);
        min_tasks = std::min(min_tasks, tasks);
        max_tasks = std::max(max_tasks, tasks);
      }

      return max_tasks > 0
                 ? static_cast<double>(max_tasks - min_tasks) / max_tasks
                 : 0.0;
    }
  };

  TbbLoadBalanceObserver() : tbb::task_scheduler_observer(false) {
    observe(true);
  }

  ~TbbLoadBalanceObserver() { observe(false); }

  void on_scheduler_entry(bool is_worker) override {
    if (is_worker) {
      // 记录线程开始工作
      thread_local size_t thread_index = GetThreadIndex();
      if (thread_index < load_metrics_.per_thread_task_counts.size()) {
        load_metrics_.per_thread_task_counts[thread_index].fetch_add(
            1, std::memory_order_relaxed);
      }
    }
  }

  // 记录任务窃取事件
  void RecordTaskSteal(bool successful) {
    if (successful) {
      load_metrics_.task_steals.fetch_add(1, std::memory_order_relaxed);
    } else {
      load_metrics_.failed_steals.fetch_add(1, std::memory_order_relaxed);
    }
  }

  // 获取负载均衡指标
  LoadBalanceMetrics GetLoadBalanceMetrics() const { return load_metrics_; }

  // 获取线程工作分布
  std::vector<uint64_t> GetThreadWorkDistribution() const {
    std::vector<uint64_t> distribution;
    distribution.reserve(load_metrics_.per_thread_task_counts.size());

    for (const auto& count : load_metrics_.per_thread_task_counts) {
      distribution.push_back(count.load(std::memory_order_acquire));
    }

    return distribution;
  }

 private:
  LoadBalanceMetrics load_metrics_;

  size_t GetThreadIndex() {
    // 获取当前线程在线程池中的索引
    // 简化实现，实际可能需要TBB内部API
    static std::atomic<size_t> thread_counter{0};
    thread_local size_t index =
        thread_counter.fetch_add(1, std::memory_order_relaxed);
    return index % tbb::info::default_concurrency();
  }
};

// 集成的TBB综合观察器
class TbbComprehensiveObserver {
 public:
  TbbComprehensiveObserver()
      : task_observer_(false), memory_observer_(), load_observer_() {}

  void StartObservation() {
    // 所有观察器已在构造函数中启动
  }

  void StopObservation() {
    // 观察器在析构函数中自动停止
  }

  // 获取综合性能报告
  struct ComprehensiveMetrics {
    TbbTaskExecutionObserver::SystemMetrics system_metrics;
    TbbMemoryObserver::MemoryMetrics memory_metrics;
    TbbMemoryObserver::SystemMemoryInfo system_memory_info;
    TbbLoadBalanceObserver::LoadBalanceMetrics load_balance_metrics;

    std::chrono::steady_clock::time_point timestamp;
  };

  ComprehensiveMetrics GetComprehensiveMetrics() const {
    return {task_observer_.GetSystemMetrics(),
            memory_observer_.GetMemoryMetrics(),
            memory_observer_.GetSystemMemoryInfo(),
            load_observer_.GetLoadBalanceMetrics(),
            std::chrono::steady_clock::now()};
  }

  // 重置所有指标
  void ResetAllMetrics() {
    task_observer_.ResetMetrics();
    // memory_observer 和 load_observer 没有reset方法，可以添加
  }

  // 设置回调函数
  void SetTaskEventCallback(
      TbbTaskExecutionObserver::TaskEventCallback callback) {
    task_observer_.SetTaskEventCallback(std::move(callback));
  }

  void SetThreadEventCallback(
      TbbTaskExecutionObserver::ThreadEventCallback callback) {
    task_observer_.SetThreadEventCallback(std::move(callback));
  }

 private:
  TbbTaskExecutionObserver task_observer_;
  TbbMemoryObserver memory_observer_;
  TbbLoadBalanceObserver load_observer_;
};
}  // namespace itbb

#endif  // ITBB_CORE_TBB_OBSERVER_H_
